import streamlit as st
import random
import json
from streamlit_tags import st_tags
import jsonlines

# Basic Streamlit page configuration
st.set_page_config(page_title='民事案件类案检索系统', layout='wide')

# Load predicate descriptions from predicate.json
with open('predicate.json', 'r', encoding='utf-8') as f:
    predicate_list = json.load(f)

# Load candidate case data from candidate.json
with open('candidate.json', 'r', encoding='utf-8') as f:
    candidate_cases = json.load(f)

# Load entity data from entity.jsonl
entity_cases = []
with jsonlines.open('entity.jsonl', 'r') as reader:
    for obj in reader:
        entity_cases.append(obj)

# Initialize session state for analysis trigger and selected case
if 'analyze_triggered' not in st.session_state:
    st.session_state['analyze_triggered'] = False
if 'selected_case' not in st.session_state:
    st.session_state['selected_case'] = None

# Page Title
# st.title('民事案件类案检索系统')

# Main Layout: Sidebar and Main Content with Spacing Adjustment
sidebar, main_content = st.columns([2, 3], gap='large')

# Sidebar Navigation
with sidebar:
    st.header('案件信息')
    case_description = st.text_area('案件描述', height=200, placeholder='请输入案件描述...')
    if st.button('分析案件'):
        if case_description.strip():
            # Randomly select a case from entity.jsonl
            st.session_state['selected_case'] = random.choice(entity_cases)
            st.session_state['analyze_triggered'] = True

# Main Content (Only show after analysis)
if st.session_state['analyze_triggered'] and st.session_state['selected_case']:
    with sidebar:
        # Extract event types and entities from the selected case
        selected_case = st.session_state['selected_case']
        event_types = list(selected_case['atom'].keys())  # List all event types, e.g., ["借款事件", ...]

        # Display detected event types
        st.subheader('检测到的法律事件类型')
        event_types
        st_tags(
            label='',
            value=event_types,
            suggestions=[],  # Disable user input
            maxtags=5,
            key='event_type_tags',
            text=""
        )
        
        # Display recognized entities for all event types, filtering out "不存在"
        st.subheader('识别出的法律实体')
        for event_type in event_types:
            entities = selected_case['atom'][event_type]  # Dictionary of entity roles and values for this event
            for role, value in entities.items():
                if value != "不存在":  # Filter out entities with value "不存在"
                    st.markdown(f"""
                    <div style='padding: 12px; background-color: #f8f9fa; border-left: 6px solid #007bff; border-radius: 8px; margin-bottom: 12px;'>
                        <strong style='color: #007bff;'>{event_type} - {role}</strong><br>
                        <span style='color: #333;'>{value}</span>
                    </div>
                    """, unsafe_allow_html=True)

        # Extracted Legal Elements (from predicate.json, keeping original logic)
        st.subheader('抽取出的法律要件')
        # Using the original sample_case legal_component for consistency
        sample_case = {
            "legal_component": [1, 0, 0, 1, 0, 0, 1]
        }
        legal_elements = [predicate_list[i] for i, v in enumerate(sample_case['legal_component']) if v == 1]
        st.markdown('<ul style="list-style-type: none; padding-left: 0;">', unsafe_allow_html=True)
        for element in legal_elements:
            st.markdown(f"<li style='padding: 10px; background-color: #ffe4e1; border-radius: 8px; margin-bottom: 8px;'><strong>{element}</strong></li>", unsafe_allow_html=True)
        st.markdown('</ul>', unsafe_allow_html=True)

    with main_content:
        st.header('类案检索结果')
        st.write('以下是与当前案件最相似的5个案例：')
        # Select random candidate cases
        candidate_keys = random.sample(list(candidate_cases.keys()), 10)
        for key in candidate_keys:
            case = candidate_cases[key]
            title = f"{case['caseID']}"
            score = round(random.uniform(0.7, 0.99), 2)  # Simulated similarity score
            legal_components = [predicate_list[i] for i, v in enumerate(case['feature']) if v == 1]
            related_components = [comp for comp in legal_components if comp in legal_elements]
            with st.expander(f"{title} (相似度: {score})"):
                st.markdown("**案例详情：**")
                st.markdown("""<div style='white-space: pre-wrap;'>
                {}
                </div>""".format("\n".join(case['content'])), unsafe_allow_html=True)
                st.markdown('<ul style="list-style-type: none; padding-left: 0;">', unsafe_allow_html=True)
                for comp in legal_components:
                    if comp in related_components:
                        st.markdown(f"<li style='padding: 10px; background-color: #ffe4e1; border-radius: 8px; margin-bottom: 8px;'><strong>{comp}</strong> (相关)</li>", unsafe_allow_html=True)
                    else:
                        st.markdown(f"<li style='padding: 10px; background-color: #e8f4f8; border-radius: 8px; margin-bottom: 8px;'>{comp}</li>", unsafe_allow_html=True)
                st.markdown('</ul>', unsafe_allow_html=True)
else:
    with main_content:
        st.info('请输入案件描述并点击 **分析案件** 以查看分析结果。')